from datetime import datetime

import torch
from torch import nn
from torch.utils.data import DataLoader
from model import GoogLetNetModel as ModelNet
from dataset import GetDataLoader

# 获取模型对象
def get_model(pth_path: str):
    # 实例化模型
    model = ModelNet()
    # 加载文件中的参数
    load_state = torch.load(pth_path)
    # 给模型加载参数
    model.load_state_dict(load_state)
    return model


# 模型测试的流程
def test_model_process(model: nn.Module, test_dataloader: DataLoader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 把模型放到设备中
    model = model.to(device)
    # 初始化参数
    test_corrects = 0.0
    # 样本的数量
    test_num = 0

    # 模型的测试，一张一张图片进行测试
    with torch.no_grad():
        # 模型测试过程中，只进行前向传播，不进行梯度计算和反向传播
        for image_x, label in test_dataloader:
            # 因为batch是1，一个批次就一张图
            image_x = image_x.to(device)
            label = label.to(device)
            # 设置模型为评估模型
            model.eval()

            # 前向传播: 10分类[...]
            output = model(image_x)
            # 取其中概率最大的值对应的索引
            pre_label = torch.argmax(output, dim=1)

            # 一张图片的预测的标签等于真实的标签，准确度加1
            test_corrects += torch.sum(pre_label == label.data)

            # 总样本数量累加
            test_num += image_x.size(0)
    # 计算准确率
    test_acc = test_corrects / test_num
    print(f"整体的{test_acc*100:.4f}%正确率，样本数：{test_num}")
    return test_acc


# 直观测试模型的流程
def test_model_process_show(model: nn.Module, test_dataloader: DataLoader):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 模型加载到设备
    model = model.to(device)
    # 标签名列表
    label_name_list = GetDataLoader().IMG_CLASSES
    # 只进行前向传播
    with torch.no_grad():
        for image_x, label in test_dataloader:
            # 图片和标签加载到设备
            image_x = image_x.to(device)
            label = label.to(device)
            # 模型设置成验证模式
            model.eval()
            # 前向传播的结果
            output = model(image_x)
            # 预测值
            pre_label = torch.argmax(output, dim=1)
            # 标签的结果
            pre_label_index = pre_label.item()
            # 真实label
            real_label_index = label.item()
            print(f'预测标签:{label_name_list[pre_label_index]},真实标签:{label_name_list[real_label_index]}')

# 测试过程，展示整体的准确率
def main():
    """
    :param path_path: 训练好的模型参数存放位置
    :return:
    """
    start = datetime.now()
    # 模型参数保存的文件
    pth_path = ModelNet.PTH_PATH
    # 获取模型对象
    model = get_model(pth_path)
    # 获取测试数据
    test_dataloader = GetDataLoader().test_dataloader
    # 测试数据，获取正确率
    test_model_process(model, test_dataloader)
    total_seconds = (datetime.now()-start).total_seconds()
    print(f'{total_seconds//60}分钟{total_seconds%60}')


# 测试过程，直接打印
def main_show():
    # 模型参数保存的文件
    pth_path = ModelNet.PTH_PATH
    # 训练好的模型参数文件路径
    model = get_model(pth_path)
    # 测试数据
    test_dataloader = GetDataLoader().test_dataloader
    # 测试过程
    test_model_process_show(model, test_dataloader)


if __name__ == "__main__":
    # 测试的整体准确率
    main()
    # 测试的每个图片的对比情况
    # main_show()